package com.artislong.config;

import com.artislong.chatmemory.MongoChatMemoryStore;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.spring.LangChain4jAutoConfig;
import dev.langchain4j.store.embedding.milvus.MilvusEmbeddingStore;
import io.milvus.common.clientenum.ConsistencyLevelEnum;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.AutoConfigureBefore;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.lang.Nullable;

import java.util.Optional;

import static com.artislong.config.MilvusEmbeddingStoreProperties.*;

/**
 * @author 陈敏
 * @version MilvusConfiguration.java, v 1.0 2025 05 19 14:31 chenmin Exp $
 * Created on 2025 05 19
 */
@AutoConfigureBefore(LangChain4jAutoConfig.class)
@AutoConfiguration
@EnableConfigurationProperties(MilvusEmbeddingStoreProperties.class)
public class MilvusEmbeddingStoreAutoConfiguration {

    @Bean
    public MilvusEmbeddingStore milvusEmbeddingStore(MilvusEmbeddingStoreProperties properties,
                                                     @Nullable EmbeddingModel embeddingModel) {
        String host = Optional.ofNullable(properties.getHost()).orElse(DEFAULT_HOST);
        int port = Optional.ofNullable(properties.getPort()).orElse(DEFAULT_PORT);
        String collectionName = Optional.ofNullable(properties.getCollectionName()).orElse(DEFAULT_COLLECTION_NAME);
        ConsistencyLevelEnum consistencyLevel = Optional.ofNullable(properties.getConsistencyLevel()).orElse(DEFAULT_CONSISTENCY_LEVEL);
        Integer dimension = Optional.ofNullable(properties.getDimension()).orElseGet(() -> embeddingModel == null ? null : embeddingModel.dimension());

        // get username and password from env variable
        String username = Optional.ofNullable(properties.getUsername()).orElse(System.getenv("MILVUS_USERNAME"));
        String password = Optional.ofNullable(properties.getPassword()).orElse(System.getenv("MILVUS_PASSWORD"));

        return MilvusEmbeddingStore.builder()
                .host(host)
                .port(port)
                .collectionName(collectionName)
                .dimension(dimension)
                .indexType(properties.getIndexType())
                .metricType(properties.getMetricType())
                .uri(properties.getUri())
                .token(properties.getToken())
                .username(username)
                .password(password)
                .consistencyLevel(consistencyLevel)
                .retrieveEmbeddingsOnSearch(properties.getRetrieveEmbeddingsOnSearch())
                .autoFlushOnInsert(properties.getAutoFlushOnInsert())
                .databaseName(properties.getDatabaseName())
                .build();
    }

    @Bean
    public ChatMemoryProvider chatMemoryProvider(MongoChatMemoryStore mongoChatMemoryStore) {
        return memoryId -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .maxMessages(10)
                .chatMemoryStore(mongoChatMemoryStore)
                .build();
    }
}

